import gym
import numpy as np
from model import Model
from Algorithm import DQN
from Agent import Agent
from ReplayMemory import ReplayMemory

LEARN_FREQ = 5  # 训练频率，不需要每一个step都learn，攒一些新增经验后再learn，提高效率
MEMORY_SIZE = 20000  # replay memory的大小，越大越占用内存
MEMORY_WARMUP_SIZE = 200  # replay_memory 里需要预存一些经验数据，再从里面sample一个batch的经验让agent去learn
BATCH_SIZE = 32  # 每次给agent learn的数据数量，从replay memory随机里sample一批数据出来
LEARNING_RATE = 0.1  # 学习率
GAMMA = 0.99  # reward 的衰减因子，一般取 0.9 到 0.999 不等


def run_episode(env, algorithm: DQN, agent, rpm):
    step = 0
    total_reward = 0
    obs = env.reset()

    while True:
        step += 1
        if len(obs) == 2:
            obs = obs[0]
        action = agent.sample(obs)
        next_obs, reward, done, _, _ = env.step(action)
        rpm.append((obs, action, reward, next_obs, done))

        if (len(rpm) > MEMORY_WARMUP_SIZE) and (step % LEARN_FREQ == 0):
            (
                batch_obs,
                batch_action,
                batch_reward,
                batch_next_obs,
                batch_done,
            ) = rpm.sample(BATCH_SIZE)
            algorithm.learn(
                batch_obs, batch_action, batch_reward, batch_next_obs, batch_done
            )

        obs = next_obs
        total_reward += reward
        if done:
            break
    return total_reward


# 评估 agent, 跑 5 个episode，总reward求平均
def evaluate(env, agent: Agent, render=False):
    eval_reward = []
    for i in range(5):
        obs = env.reset()
        episode_reward = 0
        while True:
            if len(obs) == 2:
                obs = obs[0]
            action = agent.predict(obs)  # 预测动作，只选最优动作
            obs, reward, done, _, _ = env.step(action)
            episode_reward += reward
            if render:
                env.render()
            if done:
                break
        eval_reward.append(episode_reward)
    return np.mean(eval_reward)


def main():
    env = gym.make("CartPole-v0")
    action_dim = env.action_space.n  # 2
    obs_shape = env.observation_space.shape  # (4,)

    rpm = ReplayMemory(MEMORY_SIZE)  # DQN的经验回放池
    model = Model(obs_shape[0], action_dim)
    algorithm = DQN(model, gamma=GAMMA, learnging_rate=LEARNING_RATE)
    agent = Agent(action_dim, algorithm, e_greed=0.1, e_greed_decrement=1e-6)

    # 先往经验池里存一些数据，避免最开始训练的时候样本丰富度不够
    while len(rpm) < MEMORY_WARMUP_SIZE:
        run_episode(env, algorithm, agent, rpm)

    max_episode = 2000

    # 开始训练
    episode = 0
    while episode < max_episode:  # 训练max_episode个回合，test部分不计算入episode数量
        # 训练
        for i in range(0, 50):
            total_reward = run_episode(env, algorithm, agent, rpm)
            episode += 1

        # 测试
        eval_reward = evaluate(env, agent, render=True)
        print(
            "episode:{}   e_greed:{}   Test reward:{}".format(
                episode, agent.e_greed, eval_reward
            )
        )

    # 训练结束，保存模型
    save_path = "./dqn_model.h5"
    model.model.save(save_path)
    env.close()


if __name__ == "__main__":
    main()
